{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "bfffcb27",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "2025-08-19 13:38:48.151643: I tensorflow/core/util/port.cc:153] oneDNN custom operations are on. You may see slightly different numerical results due to floating-point round-off errors from different computation orders. To turn them off, set the environment variable `TF_ENABLE_ONEDNN_OPTS=0`.\n",
      "2025-08-19 13:38:48.167235: E external/local_xla/xla/stream_executor/cuda/cuda_fft.cc:467] Unable to register cuFFT factory: Attempting to register factory for plugin cuFFT when one has already been registered\n",
      "WARNING: All log messages before absl::InitializeLog() is called are written to STDERR\n",
      "E0000 00:00:1755581928.184597 3869855 cuda_dnn.cc:8579] Unable to register cuDNN factory: Attempting to register factory for plugin cuDNN when one has already been registered\n",
      "E0000 00:00:1755581928.189692 3869855 cuda_blas.cc:1407] Unable to register cuBLAS factory: Attempting to register factory for plugin cuBLAS when one has already been registered\n",
      "W0000 00:00:1755581928.203498 3869855 computation_placer.cc:177] computation placer already registered. Please check linkage and avoid linking the same target more than once.\n",
      "W0000 00:00:1755581928.203517 3869855 computation_placer.cc:177] computation placer already registered. Please check linkage and avoid linking the same target more than once.\n",
      "W0000 00:00:1755581928.203519 3869855 computation_placer.cc:177] computation placer already registered. Please check linkage and avoid linking the same target more than once.\n",
      "W0000 00:00:1755581928.203520 3869855 computation_placer.cc:177] computation placer already registered. Please check linkage and avoid linking the same target more than once.\n",
      "2025-08-19 13:38:48.209895: I tensorflow/core/platform/cpu_feature_guard.cc:210] This TensorFlow binary is optimized to use available CPU instructions in performance-critical operations.\n",
      "To enable the following instructions: AVX2 AVX512F AVX512_VNNI FMA, in other operations, rebuild TensorFlow with the appropriate compiler flags.\n"
     ]
    }
   ],
   "source": [
    "import json\n",
    "import pandas as pd\n",
    "import os\n",
    "import numpy as np\n",
    "from rdkit import Chem\n",
    "from rdkit.Chem import AllChem, rdMolDescriptors, Descriptors, rdFreeSASA\n",
    "from sklearn.experimental import enable_iterative_imputer\n",
    "from sklearn.impute import IterativeImputer\n",
    "from sklearn.model_selection import train_test_split\n",
    "from sklearn.metrics import precision_recall_fscore_support, accuracy_score, roc_auc_score, classification_report\n",
    "import torch\n",
    "from torch.optim import Adam\n",
    "from torch_geometric.data import Dataset, Data\n",
    "from torch_geometric.loader import DataLoader\n",
    "from torch.optim.lr_scheduler import StepLR\n",
    "from torch.utils.tensorboard import SummaryWriter\n",
    "from tqdm import tqdm\n",
    "import sys\n",
    "import os\n",
    "from model import *\n",
    "import argparse\n",
    "import datetime"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "3c2e0090",
   "metadata": {},
   "outputs": [],
   "source": [
    "def predict_single(model, data, device='cpu'):\n",
    "    \"\"\"\n",
    "    对单个数据样本进行预测\n",
    "    \n",
    "    参数:\n",
    "    model (torch.nn.Module): 训练好的模型\n",
    "    data: 单个数据样本，应该包含以下属性:\n",
    "          - x: 节点特征\n",
    "          - pos: 位置信息\n",
    "          - edge_index: 边索引\n",
    "          - edge_attr: 边特征\n",
    "          - batch: 批次信息\n",
    "          - ecfp: ECFP特征\n",
    "          - maccs: MACCS特征\n",
    "          - craft_feat: craft特征\n",
    "    device (str): 设备，'cpu' 或 'cuda'\n",
    "    \n",
    "    返回:\n",
    "    dict: 包含预测结果的字典\n",
    "          - 'prediction': 预测的类别标签\n",
    "          - 'probabilities': 各类别的概率分布\n",
    "          - 'confidence': 预测的置信度（最大概率值）\n",
    "    \"\"\"\n",
    "    model.eval()\n",
    "    \n",
    "    with torch.no_grad():\n",
    "        # 将数据移动到指定设备\n",
    "        data = data.to(device)\n",
    "        \n",
    "        # 模型前向传播\n",
    "        out = model(\n",
    "            node_attr=data.x,\n",
    "            pos=data.pos,\n",
    "            edge_index=data.edge_index,\n",
    "            edge_attr=data.edge_attr,\n",
    "            batch=data.batch,\n",
    "            ecfp=data.ecfp,\n",
    "            maccs=data.maccs,\n",
    "            craft_feat=data.craft_feat\n",
    "        )\n",
    "        \n",
    "        # 计算概率分布\n",
    "        probabilities = torch.softmax(out, dim=1)\n",
    "        \n",
    "        # 获取预测标签\n",
    "        prediction = torch.argmax(probabilities, dim=1)\n",
    "        \n",
    "        # 获取置信度（最大概率值）\n",
    "        confidence = torch.max(probabilities, dim=1)[0]\n",
    "        \n",
    "        # 转换为numpy数组以便返回\n",
    "        probabilities_np = probabilities.cpu().numpy()\n",
    "        prediction_np = prediction.cpu().numpy()\n",
    "        confidence_np = confidence.cpu().numpy()\n",
    "        \n",
    "        return {\n",
    "            'prediction': prediction_np[0],  # 单个样本的预测标签\n",
    "            'probabilities': probabilities_np[0],  # 单个样本的概率分布\n",
    "            'confidence': confidence_np[0]  # 单个样本的置信度\n",
    "        }\n",
    "\n",
    "\n",
    "def predict_single_with_metrics(model, data, true_label=None, device='cpu'):\n",
    "    \"\"\"\n",
    "    对单个数据样本进行预测，并可选择性地计算与真实标签的比较\n",
    "    \n",
    "    参数:\n",
    "    model (torch.nn.Module): 训练好的模型\n",
    "    data: 单个数据样本\n",
    "    true_label (int, optional): 真实标签，如果提供将计算准确性\n",
    "    device (str): 设备，'cpu' 或 'cuda'\n",
    "    \n",
    "    返回:\n",
    "    dict: 包含预测结果和可选比较信息的字典\n",
    "    \"\"\"\n",
    "    result = predict_single(model, data, device)\n",
    "    \n",
    "    if true_label is not None:\n",
    "        result['true_label'] = true_label\n",
    "        result['is_correct'] = (result['prediction'] == true_label)\n",
    "    \n",
    "    return result"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "ec057dc3",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "模型已加载: ./model_epoch_394_best_acc_0.9887.pth\n"
     ]
    }
   ],
   "source": [
    "# 加载数据集\n",
    "train_dataset = torch.load('./dataset/train_dataset.pt', weights_only=False)\n",
    "test_dataset = torch.load('./dataset/test_dataset.pt', weights_only=False)\n",
    "train_loader = DataLoader(train_dataset, batch_size=1, shuffle=True)\n",
    "test_loader = DataLoader(test_dataset, batch_size=1, shuffle=False)\n",
    "\n",
    "device = 'cpu'\n",
    "# 模型参数\n",
    "n_layers = 1             # GNN层数\n",
    "node_dim = 7             # 节点特征维度\n",
    "out_node_nf = 227        # 输出节点特征数/分类类别数\n",
    "pos_dim = 3              # 位置编码维度\n",
    "edge_attr_dim = 3        # 边属性维度\n",
    "hidden_f = 128           # 隐藏层维度\n",
    "craft_dim = 29           # CRAFT特征维度\n",
    "craft_hidden = 58        # CRAFT隐藏层维度\n",
    "craft_out_feat = 29      # CRAFT输出特征数\n",
    "\n",
    "model = ZEOEGNN4(\n",
    "    n_layers=n_layers,\n",
    "    node_dim=node_dim,\n",
    "    out_node_nf=out_node_nf,\n",
    "    pos_dim=pos_dim,\n",
    "    edge_attr_dim=edge_attr_dim,\n",
    "    hidden_f=hidden_f,\n",
    "    craft_dim=craft_dim,\n",
    "    craft_hidden=craft_hidden,\n",
    "    craft_out_feat=craft_out_feat\n",
    ")\n",
    "\n",
    "model_path = './model_epoch_394_best_acc_0.9887.pth'\n",
    "if os.path.exists(model_path):\n",
    "    model.load_state_dict(torch.load(model_path, map_location=device))\n",
    "    print(f\"模型已加载: {model_path}\")\n",
    "else:\n",
    "    print(f\"模型文件 {model_path} 不存在，请检查路径。\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "b8d1675c",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Batch 1 - 预测类别: 21 / 21, 置信度: 1.0000\n",
      "Batch 2 - 预测类别: 7 / 7, 置信度: 1.0000\n",
      "Batch 3 - 预测类别: 98 / 81, 置信度: 1.0000\n",
      "Batch 4 - 预测类别: 7 / 101, 置信度: 1.0000\n",
      "Batch 5 - 预测类别: 7 / 7, 置信度: 1.0000\n",
      "Batch 6 - 预测类别: 12 / 12, 置信度: 1.0000\n",
      "Batch 7 - 预测类别: 7 / 7, 置信度: 1.0000\n",
      "Batch 8 - 预测类别: 75 / 75, 置信度: 1.0000\n",
      "Batch 9 - 预测类别: 76 / 76, 置信度: 1.0000\n",
      "Batch 10 - 预测类别: 43 / 43, 置信度: 1.0000\n",
      "Batch 11 - 预测类别: 84 / 43, 置信度: 1.0000\n",
      "Batch 12 - 预测类别: 98 / 4, 置信度: 1.0000\n",
      "Batch 13 - 预测类别: 0 / 0, 置信度: 1.0000\n",
      "Batch 14 - 预测类别: 0 / 0, 置信度: 1.0000\n",
      "Batch 15 - 预测类别: 98 / 176, 置信度: 1.0000\n",
      "Batch 16 - 预测类别: 45 / 45, 置信度: 1.0000\n",
      "Batch 17 - 预测类别: 75 / 75, 置信度: 1.0000\n",
      "Batch 18 - 预测类别: 84 / 43, 置信度: 1.0000\n",
      "Batch 19 - 预测类别: 7 / 7, 置信度: 1.0000\n",
      "Batch 20 - 预测类别: 71 / 16, 置信度: 1.0000\n",
      "Batch 21 - 预测类别: 8 / 8, 置信度: 1.0000\n"
     ]
    }
   ],
   "source": [
    "# 循环加载batch\n",
    "for i, data in enumerate(test_loader):\n",
    "    result = predict_single(model, data)\n",
    "    print(f\"Batch {i+1} - 预测类别: {result['prediction']} / {data.y.item()}, 置信度: {result['confidence']:.4f}\")\n",
    "    if i >= 20:  # 只预测前10个batch\n",
    "        break"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "41db46a3",
   "metadata": {},
   "source": [
    "#### 创建一个随机数据用于测试"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "ab5d8109",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Batch 21 - 预测类别: 7 / 8, 置信度: 1.0000\n"
     ]
    }
   ],
   "source": [
    "import random\n",
    "def generate_random_sample(n1: int, n2: int) -> dict:\n",
    "    \"\"\"\n",
    "    生成随机样本数据，符合 PredictionRequest 数据模型。\n",
    "\n",
    "    参数:\n",
    "    n1 (int): 分子原子数（节点数），例如 37\n",
    "    n2 (int): 边数量，例如 76\n",
    "\n",
    "    返回:\n",
    "    dict: 包含 PredictionRequest 所需字段的字典，数据类型和形状符合要求\n",
    "\n",
    "    Raises:\n",
    "    ValueError: 如果 n1 <= 0 或 n2 < 0 或 n2 过大导致无法生成有效边\n",
    "    \"\"\"\n",
    "    # 输入验证\n",
    "    if n1 <= 0:\n",
    "        raise ValueError(\"节点数 n1 必须为正整数\")\n",
    "    if n2 < 0:\n",
    "        raise ValueError(\"边数 n2 必须为非负整数\")\n",
    "    max_edges = n1 * (n1 - 1) // 2  # 无向图的最大边数\n",
    "    if n2 > max_edges:\n",
    "        raise ValueError(f\"边数 n2 ({n2}) 不能超过无向图最大边数 {max_edges}\")\n",
    "\n",
    "    # 1. 节点特征 (x): [n1, 7], torch.float32\n",
    "    x = torch.randn(n1, 7, dtype=torch.float32).tolist()\n",
    "\n",
    "    # 2. 边索引 (edge_index): [2, n2], torch.int64\n",
    "    # 生成有效边，确保边连接的节点索引在 [0, n1-1] 范围内\n",
    "    edge_index = []\n",
    "    if n2 > 0:\n",
    "        # 随机生成 n2 条边，避免自环和重复边\n",
    "        possible_edges = [(i, j) for i in range(n1) for j in range(i + 1, n1)]\n",
    "        random.shuffle(possible_edges)\n",
    "        selected_edges = possible_edges[:min(n2, len(possible_edges))]\n",
    "        edge_index = [[e[0], e[1]] for e in selected_edges] + [[e[1], e[0]] for e in selected_edges]  # 无向图，双向边\n",
    "        edge_index = edge_index[:n2]  # 截取精确的 n2 条边\n",
    "        edge_index = [list(e) for e in zip(*edge_index)]  # 转换为 [2, n2] 格式\n",
    "        if len(edge_index[0]) < n2:  # 如果边数不足，补充随机边\n",
    "            remaining = n2 - len(edge_index[0])\n",
    "            extra_edges = torch.randint(0, n1, (2, remaining), dtype=torch.int64).tolist()\n",
    "            edge_index[0].extend(extra_edges[0])\n",
    "            edge_index[1].extend(extra_edges[1])\n",
    "    else:\n",
    "        edge_index = [[], []]  # 空图情况\n",
    "\n",
    "    # 3. 边属性 (edge_attr): [n2, 3], torch.float32\n",
    "    edge_attr = torch.randn(n2, 3, dtype=torch.float32).tolist()\n",
    "\n",
    "    # 4. 节点位置 (pos): [n1, 3], torch.float32\n",
    "    pos = torch.randn(n1, 3, dtype=torch.float32).tolist()\n",
    "\n",
    "    # 5. 节点属性 (z): [n1], torch.int64\n",
    "    # 假设 z 表示原子序数，随机生成 1-100 之间的整数（常见元素范围）\n",
    "    z = torch.randint(1, 101, (n1,), dtype=torch.int64).tolist()\n",
    "\n",
    "    # 6. SMILES 字符串: str\n",
    "    # 使用简单的占位符，长度随机，实际应用中应使用真实 SMILES\n",
    "    smiles = \"C\" * random.randint(5, 20)\n",
    "\n",
    "    # 7. MACCS 指纹: [1, 167], torch.int64 (二进制)\n",
    "    maccs = torch.randint(0, 2, (1, 167), dtype=torch.float32).tolist()\n",
    "\n",
    "    # 8. ECFP 指纹: [1, 2048], torch.int64 (二进制)\n",
    "    ecfp = torch.randint(0, 2, (1, 2048), dtype=torch.float32).tolist()\n",
    "\n",
    "    # 9. 标签 (y): int\n",
    "    y = random.randint(0, 1)  # 假设二分类问题，标签为 0 或 1\n",
    "\n",
    "    # 10. CRAFT 特征: [1, 29], torch.float32\n",
    "    craft_feat = torch.randn(1, 29, dtype=torch.float32).tolist()\n",
    "\n",
    "    # 返回符合 PredictionRequest 的字典\n",
    "    return {\n",
    "        \"x\": x,\n",
    "        \"edge_index\": edge_index,\n",
    "        \"edge_attr\": edge_attr,\n",
    "        \"pos\": pos,\n",
    "        \"z\": z,\n",
    "        \"smiles\": smiles,\n",
    "        \"maccs\": maccs,\n",
    "        \"ecfp\": ecfp,\n",
    "        \"y\": y,\n",
    "        \"craft_feat\": craft_feat\n",
    "    }\n",
    "# 生成一个随机样本，n1=10（节点数），n2=12（边数）\n",
    "sample = generate_random_sample(n1=37, n2=76)\n",
    "\n",
    "data_obj = Data(\n",
    "    x=torch.tensor(sample[\"x\"], dtype=torch.float32),\n",
    "    edge_index=torch.tensor(sample[\"edge_index\"], dtype=torch.int64),\n",
    "    edge_attr=torch.tensor(sample[\"edge_attr\"], dtype=torch.float32),\n",
    "    pos=torch.tensor(sample[\"pos\"], dtype=torch.float32),\n",
    "    z=torch.tensor(sample[\"z\"], dtype=torch.int64),\n",
    "    smiles=sample[\"smiles\"],\n",
    "    maccs=torch.tensor(sample[\"maccs\"], dtype=torch.float32),\n",
    "    ecfp=torch.tensor(sample[\"ecfp\"], dtype=torch.float32),\n",
    "    y=sample[\"y\"],\n",
    "    craft_feat=torch.tensor(sample[\"craft_feat\"], dtype=torch.float32)\n",
    ")\n",
    "\n",
    "\n",
    "result = predict_single(model, data_obj)\n",
    "print(f\"Batch {i+1} - 预测类别: {result['prediction']} / {data.y.item()}, 置信度: {result['confidence']:.4f}\")\n",
    "# result"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 43,
   "id": "8d3ccf1f",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "x: dtype=torch.float32, shape=torch.Size([37, 7])\n",
      "edge_index: dtype=torch.int64, shape=torch.Size([2, 76])\n",
      "edge_attr: dtype=torch.float32, shape=torch.Size([76, 3])\n",
      "pos: dtype=torch.float32, shape=torch.Size([37, 3])\n",
      "z: dtype=torch.int64, shape=torch.Size([37])\n",
      "smiles: dtype=str, shape=N/A\n",
      "maccs: dtype=torch.int64, shape=torch.Size([1, 167])\n",
      "ecfp: dtype=torch.int64, shape=torch.Size([1, 2048])\n",
      "y: dtype=int, shape=N/A\n",
      "craft_feat: dtype=torch.float32, shape=torch.Size([1, 29])\n"
     ]
    }
   ],
   "source": [
    "# 假设 data_obj 是一个 torch_geometric.data.Data 对象\n",
    "for attr in ['x', 'edge_index', 'edge_attr', 'pos', 'z', 'smiles', 'maccs', 'ecfp', 'y', 'craft_feat']:\n",
    "    value = getattr(data_obj, attr)\n",
    "    if torch.is_tensor(value):\n",
    "        dtype = value.dtype\n",
    "        shape = value.shape\n",
    "    else:\n",
    "        dtype = type(value).__name__  # 获取 Python 类型名（如 'str' 或 'int'）\n",
    "        shape = 'N/A'\n",
    "    print(f\"{attr}: dtype={dtype}, shape={shape}\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 48,
   "id": "ef5f09f9",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "x: dtype=torch.float32, shape=torch.Size([37, 7])\n",
      "edge_index: dtype=torch.int64, shape=torch.Size([2, 76])\n",
      "edge_attr: dtype=torch.float32, shape=torch.Size([76, 3])\n",
      "pos: dtype=torch.float32, shape=torch.Size([37, 3])\n",
      "z: dtype=torch.int64, shape=torch.Size([37])\n",
      "smiles: dtype=str, shape=N/A\n",
      "maccs: dtype=torch.float32, shape=torch.Size([1, 167])\n",
      "ecfp: dtype=torch.float32, shape=torch.Size([1, 2048])\n",
      "y: dtype=int, shape=N/A\n",
      "craft_feat: dtype=torch.float32, shape=torch.Size([1, 29])\n"
     ]
    }
   ],
   "source": [
    "test_loader.dataset[15]\n",
    "for attr in ['x', 'edge_index', 'edge_attr', 'pos', 'z', 'smiles', 'maccs', 'ecfp', 'y', 'craft_feat']:\n",
    "    value = getattr(test_loader.dataset[15], attr)\n",
    "    if torch.is_tensor(value):\n",
    "        dtype = value.dtype\n",
    "        shape = value.shape\n",
    "    else:\n",
    "        dtype = type(value).__name__  # 获取 Python 类型名（如 'str' 或 'int'）\n",
    "        shape = 'N/A'\n",
    "    print(f\"{attr}: dtype={dtype}, shape={shape}\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4057b4ca",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "Data(x=[37, 7], edge_index=[2, 76], edge_attr=[76, 3], y=3, pos=[37, 3], z=[37], smiles='CCCCCCCCCCCCCCCCC', maccs=[1, 167], ecfp=[1, 2048], craft_feat=[1, 29])"
      ]
     },
     "execution_count": 29,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "data_obj"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "gnn_0",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.10.16"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
